package com.yricky.atri.viewmodel

import android.graphics.Bitmap
import com.yricky.atri.utils.yolov5.Recognition
import com.yricky.mvi.StateViewModel
import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.gpu.GpuDelegate
import org.tensorflow.lite.nnapi.NnApiDelegate

/**
 * @author Yricky
 * @date 2022/1/9
 */
class YoloV5SessionViewModel:StateViewModel<YoloV5SessionViewModel.ViewState>() {
    data class ViewState(
        val results:List<Recognition> = emptyList(),
        val originalBitmap: Bitmap = Bitmap.createBitmap(1,1,Bitmap.Config.ARGB_8888),
        val avgClassifyMs:Long = 0,
        val avgInterpreterMs:Long = 0,
        val isPrepared:Boolean = false,
        val useBackCam:Boolean = true,
        val useDevice: UseDevice = UseDevice.CPU
    )

    sealed class UseDevice(val option:Interpreter.Options){

        object CPU:UseDevice(Interpreter.Options())
        object GPU:UseDevice(Interpreter.Options().addDelegate(GpuDelegate()))
        object NNApi:UseDevice(Interpreter.Options().addDelegate(NnApiDelegate(
            NnApiDelegate.Options().setAllowFp16(true)
        )))
    }

    fun changeDevice(dev:UseDevice){
        if(dev != viewState.value.useDevice)
            updateState { it.copy(useDevice = dev) }
    }

    override val initState: () -> ViewState = { ViewState() }
}